import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import torch
import torchvision.transforms as transforms
from PIL import Image, ImageDraw, ImageOps
import numpy as np
import threading
import matplotlib.pyplot as plt
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import sys
import os
import matplotlib as mpl
from matplotlib import font_manager

# 尝试加载中文字体
try:
    font_path = 'font/siyuan.otf' 
    if os.path.exists(font_path):
        font_manager.fontManager.addfont(font_path)
        plt.rcParams['font.family'] = font_manager.FontProperties(fname=font_path).get_name()
    else:
        print("字体文件不存在，使用默认字体")
except Exception as e:
    print(f"字体加载失败: {e}，使用默认字体")

from data.data_loader import MNISTLoader
from model.cnn import CNN
from utils.trainer import Trainer
from utils.visualize import TrainingVisualizer, ActivationVisualizer

class MNISTApp:
    def __init__(self, root):
        """初始化应用程序"""
        self.root = root
        self.root.title("MNIST手写数字识别")
        self.root.geometry("1400x900")
        self.root.minsize(1200, 800)
        
        # 初始化模型和训练器
        self.model = None
        self.trainer = None
        self.visualizer = None
        self.activation_visualizer = None
        self.train_loader = None
        self.test_loader = None
        
        # 线程控制标志
        self.running = True
        self._cleaning_up = False
        self.training_cancelled = False
        self.train_thread = None
        self.recognize_threads = []
        
        # 预测相关参数 
        self.top_n_predictions = 3
        
        # 设置设备
        self.device = self.get_best_device()
        
        # 训练配置
        self.train_config = {
            'batch_size': 64,
            'learning_rate': 0.001,
            'epochs': 10,
            'weight_decay': 1e-4,
            'betas': (0.9, 0.999),
            'early_stopping': True,
            'patience': 10,
            'gradient_clip': 1.0,
            'use_augmentation': True,
            'early_stopping_patience': 5,
            'lr_scheduler': 'onecycle',
            'augmentation_level': 'medium',
            'scheduler_params': {
                'onecycle': {'max_lr_multiplier': 25.0, 'pct_start': 0.3},
                'cosine': {'T_max': 10, 'min_lr_multiplier': 0.01},
                'step': {'step_size': 5, 'gamma': 0.5}
            },
            'top_n_predictions': 3  # 添加预测结果显示数量
        }
        
        # 训练参数
        self.train_params = {
            'learning_rate': 0.001,
            'batch_size': 64,
            'epochs': 10,
            'leaky_relu_alpha': 0.01,
            'num_conv_layers': 3,
            'base_channels': 32,
            'dropout_rate': 0.3,
            'use_batch_norm': True,
            'use_gmp': False,
            'use_mlp': True,
            'mlp_hidden_dims': (256, 128),
            'num_classes': 10
        }
        
        # 数据加载
        self.data_loader = MNISTLoader()
        
        # 初始化训练历史数据
        self.epochs = []
        self.train_losses = []
        self.test_losses = []
        self.train_accs = []
        self.test_accs = []
        self.lr_history = []
        
        # 创建界面
        self.create_widgets()
        
        # 加载数据
        self.load_data()
        
    def get_best_device(self):
        # 检查CUDA
        if torch.cuda.is_available():
            device = torch.device("cuda")
            gpu_info = torch.cuda.get_device_properties(0)
            print(f"使用GPU: {gpu_info.name}, 内存: {gpu_info.total_memory / 1024**3:.2f}GB")
        # 检查MPS (Apple Silicon)
        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            device = torch.device("mps")
            print("使用Apple MPS设备")
        else:
            device = torch.device("cpu")
            print("使用CPU设备")
        return device
    
    def create_widgets(self):
        # 创建主框架
        main_frame = ttk.Frame(self.root, padding="10")
        main_frame.pack(fill=tk.BOTH, expand=True)
        
        # 创建左侧控制面板
        control_frame = ttk.LabelFrame(main_frame, text="控制面板", padding="10")
        control_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=False, padx=5, pady=5)
        
        # 数据加载按钮
        ttk.Button(control_frame, text="加载数据", command=self.load_data).pack(fill=tk.X, pady=2)
        
        # 参数设置按钮
        ttk.Button(control_frame, text="设置参数", command=self.show_train_params_dialog).pack(fill=tk.X, pady=2)
        
        # 训练按钮
        self.train_btn = ttk.Button(control_frame, text="训练模型", command=self.start_training)
        self.train_btn.pack(fill=tk.X, pady=2)
        
        # 取消训练按钮
        self.cancel_train_btn = ttk.Button(control_frame, text="取消训练", command=self.cancel_training, state=tk.DISABLED)
        self.cancel_train_btn.pack(fill=tk.X, pady=2)
        
        # 进度条
        self.progress = ttk.Progressbar(control_frame, orient=tk.HORIZONTAL, length=100, mode='determinate')
        self.progress.pack(fill=tk.X, pady=5)
        
        # 模型保存按钮
        ttk.Button(control_frame, text="保存模型", command=self.save_model).pack(fill=tk.X, pady=2)
        
        # 模型加载按钮
        ttk.Button(control_frame, text="加载模型", command=self.load_model).pack(fill=tk.X, pady=2)
        
        # 激活值可视化按钮
        ttk.Button(control_frame, text="查看激活值", command=self.show_activations).pack(fill=tk.X, pady=2)
        
        # === 手写识别控制区域 ===
        draw_control_frame = ttk.LabelFrame(control_frame, text="手写识别控制", padding="5")
        draw_control_frame.pack(fill=tk.X, pady=10)
        
        # 按钮框架
        btn_frame = ttk.Frame(draw_control_frame)
        btn_frame.pack(fill=tk.X, pady=5)
        
        # 清除按钮
        ttk.Button(btn_frame, text="清除画布", command=self.clear_canvas).pack(side=tk.LEFT, padx=2, fill=tk.X, expand=True)
        
        # 识别按钮
        ttk.Button(btn_frame, text="识别数字", command=self.recognize_digit).pack(side=tk.LEFT, padx=2, fill=tk.X, expand=True)
        
        # 结果显示区域
        result_label_frame = ttk.LabelFrame(draw_control_frame, text="预测结果", padding="5")
        result_label_frame.pack(fill=tk.X, pady=5)

        # 主要预测结果
        main_result_frame = ttk.Frame(result_label_frame)
        main_result_frame.pack(fill=tk.X, pady=2)

        self.result_var = tk.StringVar(value="未识别")
        ttk.Label(main_result_frame, text="识别结果:", font=('Arial', 10)).pack(side=tk.LEFT)
        ttk.Label(main_result_frame, textvariable=self.result_var, font=('Arial', 12, 'bold')).pack(side=tk.LEFT, padx=(5,0))

        # 置信度
        confidence_frame = ttk.Frame(result_label_frame)
        confidence_frame.pack(fill=tk.X, pady=2)

        self.confidence_var = tk.StringVar(value="0%")
        ttk.Label(confidence_frame, text="置信度:", font=('Arial', 10)).pack(side=tk.LEFT)
        ttk.Label(confidence_frame, textvariable=self.confidence_var, font=('Arial', 10)).pack(side=tk.LEFT, padx=(5,0))
        
        # 前N个预测结果显示区域
        self.top_predictions_frame = ttk.Frame(result_label_frame)
        self.top_predictions_frame.pack(fill=tk.X, pady=5)
        
        # 初始化空的预测结果显示
        ttk.Label(self.top_predictions_frame, text="等待识别...", font=('Arial', 10, 'italic')).pack(anchor=tk.W)
        
        # 设备信息显示
        device_info_frame = ttk.LabelFrame(control_frame, text="设备信息", padding="5")
        device_info_frame.pack(fill=tk.X, pady=10)
        
        device_text = f"当前设备: {self.device}"
        if torch.cuda.is_available():
            gpu_info = torch.cuda.get_device_properties(0)
            device_text += f"\nGPU: {gpu_info.name}\n内存: {gpu_info.total_memory / 1024**3:.2f}GB"
        
        ttk.Label(device_info_frame, text=device_text).pack(anchor=tk.W)
        
        # 创建右侧区域的主框架
        right_main_frame = ttk.Frame(main_frame)
        right_main_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=5, pady=5)
        

        draw_frame = ttk.LabelFrame(right_main_frame, text="手写区域 (280x280)", padding="10")
        draw_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=False, padx=5, pady=5)
        
        # 固定尺寸的画布，确保280x280px
        self.canvas = tk.Canvas(draw_frame, width=280, height=280, bg="white", highlightthickness=1, highlightbackground="black")
        self.canvas.pack(pady=10)
        
        # 绑定鼠标事件
        self.canvas.bind("<B1-Motion>", self.paint)
        self.canvas.bind("<ButtonRelease-1>", self.reset)
        

        viz_frame = ttk.LabelFrame(right_main_frame, text="训练可视化", padding="10")
        viz_frame.pack(side=tk.BOTTOM, fill=tk.BOTH, expand=True, padx=5, pady=5)
        
        # 创建左右分栏的可视化框架
        viz_container = ttk.Frame(viz_frame)
        viz_container.pack(fill=tk.BOTH, expand=True)
        
        # 左侧图表框架 - 损失曲线
        loss_frame = ttk.LabelFrame(viz_container, text="损失曲线")
        loss_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=2, pady=2, ipadx=10, ipady=10)
        
        # 右侧图表框架 - 准确率曲线
        acc_frame = ttk.LabelFrame(viz_container, text="准确率曲线")
        acc_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=2, pady=2, ipadx=10, ipady=10)
        
        # 创建两个独立的可视化图表，确保大小一致
        figsize = (5, 4)
        self.loss_visualizer = TrainingVisualizer(
            parent_frame=loss_frame,
            use_tkagg=True,
            figsize=figsize,
            plot_type='loss'
        )
        
        self.acc_visualizer = TrainingVisualizer(
            parent_frame=acc_frame,
            use_tkagg=True,
            figsize=figsize,
            plot_type='accuracy'
        )
        
        # 存储绘图数据
        self.last_x = None
        self.last_y = None
        self.line_width = 8
    
    def load_data(self):
        """加载MNIST数据，使用数据加载器"""
        try:
            # 根据配置决定是否使用数据增强
            use_augmentation = self.train_config.get('use_augmentation', True)
            
            # 加载数据集，传递use_augmentation参数来控制是否启用数据增强
            self.data_loader.load_datasets(
                validation_split=0.1,
                advanced_augmentation=use_augmentation and self.train_config['augmentation_level'] in ['medium', 'heavy'],
                use_augmentation=use_augmentation
            )
            
            # 创建数据加载器
            self.train_loader, self.val_loader, self.test_loader = self.data_loader.create_dataloaders(
                batch_size=self.train_config['batch_size'],
                num_workers=4,
                shuffle_train=True,
                pin_memory=True,
                persistent_workers=True
            )
            
            # 分析数据分布
            distribution = self.data_loader.get_class_distribution()
            print(f"训练集类别分布: {distribution['train']}")
            print(f"验证集类别分布: {distribution['val']}")
            print(f"测试集类别分布: {distribution['test']}")
            
            # 重置训练可视化
            if hasattr(self, 'loss_visualizer') and self.loss_visualizer:
                self.loss_visualizer.clear()
                self.acc_visualizer.clear()
            
            # 初始化绘图数据
            self.epochs = []
            self.train_losses = []
            self.test_losses = []
            self.train_accs = []
            self.test_accs = []
            self.lr_history = []
            
            messagebox.showinfo("成功", "MNIST数据加载成功!")
        except Exception as e:
            messagebox.showerror("错误", f"数据加载失败: {str(e)}")
    
    def paint(self, event):
        """在画布上绘制"""
        x, y = event.x, event.y
        if self.last_x and self.last_y:
            self.canvas.create_line(self.last_x, self.last_y, x, y, 
                                   width=self.line_width, fill="black", 
                                   capstyle=tk.ROUND, smooth=tk.TRUE)
        self.last_x = x
        self.last_y = y
    
    def reset(self, event):
        """重置鼠标位置"""
        self.last_x = None
        self.last_y = None
    
    def clear_canvas(self):
        """清除画布"""
        self.canvas.delete("all")
        self.result_var.set("未识别")
        self.confidence_var.set("0%")
        
        # 重置前N个预测结果显示
        for widget in self.top_predictions_frame.winfo_children():
            widget.destroy()
        ttk.Label(self.top_predictions_frame, text="等待识别...", font=('Arial', 10, 'italic')).pack(anchor=tk.W)
        
        # 重置绘图状态变量，确保画布清空后能正确重绘
        self.last_x = None
        self.last_y = None
    
    def get_canvas_image(self):
        """从画布获取图像并转换为MNIST格式"""
        # 创建一个空白的RGBA图像
        width = 280
        height = 280
        img = Image.new('RGB', (width, height), color='white')
        draw = ImageDraw.Draw(img)
        
        # 获取画布上的所有项目
        items = self.canvas.find_all()
        
        # 绘制线条
        for item in items:
            # 获取项目类型
            item_type = self.canvas.type(item)
            
            if item_type == 'line':
                # 获取线条坐标
                coords = self.canvas.coords(item)
                # 获取线条属性
                width = self.canvas.itemcget(item, 'width')
                fill = self.canvas.itemcget(item, 'fill')
                
                # 绘制线条
                draw.line(coords, fill=fill, width=int(float(width)))
        
        # 转换为灰度图像
        img = img.convert('L')
        
        # 调整大小为28x28
        img = img.resize((28, 28), Image.LANCZOS)
        
        # 反色（MNIST是黑底白字）
        img = ImageOps.invert(img)
        
        # 转换为numpy数组
        img_array = np.array(img)
        
        img_array = img_array.astype(np.float32)
        
        # 应用图像中心化，确保数字居中
        non_zero = np.where(img_array > 10) 
        if len(non_zero[0]) > 0:
            # 计算边界框
            min_y, max_y = non_zero[0].min(), non_zero[0].max()
            min_x, max_x = non_zero[1].min(), non_zero[1].max()
            
            # 计算数字大小
            h = max_y - min_y + 1
            w = max_x - min_x + 1
            
            # 创建一个新的28x28图像，将数字居中
            centered_img = np.zeros((28, 28), dtype=np.float32)
            
            # 计算居中位置
            offset_y = max(0, (28 - h) // 2)
            offset_x = max(0, (28 - w) // 2)
            
            # 确保不会越界
            h = min(h, 28 - offset_y)
            w = min(w, 28 - offset_x)
            
            # 复制数字到居中位置
            centered_img[offset_y:offset_y + h, offset_x:offset_x + w] = img_array[min_y:min_y + h, min_x:min_x + w]
            img_array = centered_img
        

        # 归一化
        img_array = img_array / 255.0
        
        # 转换为PyTorch张量
        img_tensor = torch.tensor(img_array, dtype=torch.float32)
        
        # 添加批次维度和通道维度，确保形状为 [1, 1, 28, 28]
        img_tensor = img_tensor.unsqueeze(0).unsqueeze(0)
        

        transform = transforms.Compose([
            transforms.Normalize(mean=[0.1307], std=[0.3081])
        ])
        img_tensor = transform(img_tensor)
        
        # 打印详细调试信息
        print("[调试] 处理后的张量形状:", img_tensor.shape)
        print("[调试] 张量统计信息:")
        print(f"  均值: {img_tensor.mean().item():.4f}")
        print(f"  标准差: {img_tensor.std().item():.4f}")
        print(f"  最大值: {img_tensor.max().item():.4f}")
        print(f"  最小值: {img_tensor.min().item():.4f}")
        
        return img_tensor
    
    def recognize_digit(self):
        """识别手写数字，支持显示前N个预测结果"""
        if self.model is None:
            messagebox.showwarning("警告", "请先训练或加载模型!")
            return
        

        top_n = self.train_config.get('top_n_predictions', 3)
        
        # 在单独线程中进行识别，避免UI卡顿
        def recognize_thread():
            try:
                # 检查是否应该继续执行
                if not self.running:
                    return
                
                # 获取图像
                img_tensor = self.get_canvas_image()
                
                # 再次检查
                if not self.running:
                    return
                
                # 进行预测
                prediction, probabilities = self.trainer.predict(img_tensor)
                digit = prediction.item()
                confidence = probabilities[0][digit].item() * 100
                
                # 获取前N个预测结果
                probs_np = probabilities.cpu().numpy()[0]
                top_indices = np.argsort(probs_np)[::-1][:top_n]
                top_probs = probs_np[top_indices]
                
                # 打印所有类别的概率分布，帮助诊断问题
                print("[调试] 所有数字的预测概率:")
                for i in range(10):
                    print(f"  数字{i}: {probs_np[i]*100:.2f}%")
                
                # 计算概率分布的熵，用于评估模型的确定性
                entropy = -np.sum(probs_np * np.log(probs_np + 1e-10))
                print(f"[调试] 概率分布熵: {entropy:.4f} (熵越低，确定性越高)")
                
                # 再次检查
                if not self.running:
                    return
                
                # 在主线程中更新结果
                self.root.after(0, lambda: self.result_var.set(str(digit)))
                self.root.after(0, lambda: self.confidence_var.set(f"{confidence:.2f}%"))
                
                # 更新前N个预测结果显示
                def update_top_predictions():
                    try:
                        # 确保top_predictions_frame存在
                        if not hasattr(self, 'top_predictions_frame') or not self.top_predictions_frame:
                            print("警告: top_predictions_frame未正确初始化")
                            return
                        
                        # 清除之前的预测结果
                        for widget in self.top_predictions_frame.winfo_children():
                            widget.destroy()
                        
                        # 打印调试信息
                        print(f"预测结果: 数字{digit}，置信度{confidence:.2f}%")
                        print(f"前{top_n}个预测结果: {list(zip(top_indices, top_probs*100))}")
                        
                        # 显示前N个预测结果
                        for i in range(min(top_n, len(top_indices))):
                            idx = top_indices[i]
                            prob = top_probs[i]
                            result_text = f"{i+1}. 数字 {idx} ({prob*100:.2f}%)"
                            if idx == digit:  # 最高分结果使用不同样式
                                ttk.Label(self.top_predictions_frame, text=result_text, 
                                         font=('Arial', 10, 'bold'), foreground='#0066cc').pack(anchor=tk.W)
                            else:
                                ttk.Label(self.top_predictions_frame, text=result_text, 
                                         font=('Arial', 10)).pack(anchor=tk.W)
                    except Exception as e:
                        print(f"更新预测结果时出错: {str(e)}")
                        # 显示错误信息到界面
                        ttk.Label(self.top_predictions_frame, text=f"错误: {str(e)}", 
                                 font=('Arial', 10), foreground='red').pack(anchor=tk.W)
                
                self.root.after(0, update_top_predictions)
                
            except Exception as e:
                # 错误消息在主线程中显示，使用默认参数捕获异常信息
                if self.running:
                    self.root.after(0, lambda err_msg=str(e): messagebox.showerror("错误", f"识别失败: {err_msg}"))
        
        # 启动识别线程
        thread = threading.Thread(target=recognize_thread)
        self.recognize_threads.append(thread)
        thread.start()
        
        # 清理已完成的线程
        self.recognize_threads = [t for t in self.recognize_threads if t.is_alive()]
    
    def show_activations(self):
        """
        显示激活值可视化窗口
        """
        if not hasattr(self, 'model') or self.model is None:
            messagebox.showerror("错误", "请先训练或加载模型！")
            return
        
        # 检查画布是否有内容
        if not hasattr(self, 'canvas'):
            messagebox.showinfo("提示", "画布未初始化！")
            return
        
        # 尝试从画布获取图像数据
        try:
            # 获取画布的图像数据
            self.image_data = self.canvas.postscript(colormode='mono')
            if not self.image_data or len(self.image_data) < 100:
                messagebox.showinfo("提示", "请先在画布上绘制数字！")
                return
        except Exception:
            messagebox.showinfo("提示", "无法获取画布内容，请先绘制数字！")
            return
        
        # 创建新窗口
        self.activation_window = tk.Toplevel(self.root)
        
        # 创建容纳画布的框架
        frame = ttk.Frame(self.activation_window)
        frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        try:
            # 使用现有的get_canvas_image方法获取处理后的图像张量
            img_tensor = self.get_canvas_image()
            
            # 打印调试信息 - 检查图像数据
            print("[调试] 图像张量形状:", img_tensor.shape)
            print("[调试] 图像张量统计:", "均值=", img_tensor.mean().item(), 
                  "最大值=", img_tensor.max().item(), "最小值=", img_tensor.min().item())
            
            # 确保trainer存在
            if not hasattr(self, 'trainer') or self.trainer is None:
                messagebox.showerror("错误", "训练器未初始化！")
                return
            
            # 预先检查模型预测结果
            with torch.no_grad():
                test_input = img_tensor.to(self.device)
                test_output = self.model(test_input)
                test_prob = torch.softmax(test_output, dim=1)
                _, test_pred = torch.max(test_output, 1)
                print(f"[调试] 测试预测结果: 数字{test_pred.item()}，置信度{test_prob[0][test_pred].item()*100:.2f}%")
            
            # 创建支持Tkinter的激活值可视化器
            self.activation_visualizer = ActivationVisualizer(
                max_channels=6, 
                cmap='viridis', 
                show_colorbar=True,
                parent_frame=frame,
                use_tkagg=True
            )
            
            # 获取模型各层的激活值
            activations = self.trainer.get_activations(img_tensor)
            
            # 确保activations是字典
            if not isinstance(activations, dict):
                # 如果返回的不是字典，尝试使用默认层名称
                print("[调试] activations不是字典，使用默认层名称")
                num_layers = len(activations)
                default_names = [f'层{i+1}' for i in range(num_layers)]
                self.activation_visualizer.update(img_tensor, activations, layer_names=default_names)
            else:
                # 正常更新可视化
                self.activation_visualizer.update(img_tensor, list(activations.values()), layer_names=list(activations.keys()))
        except Exception as e:
            messagebox.showerror("错误", f"获取激活值失败: {str(e)}")
        
        # 窗口
        self.activation_window.title("激活值可视化")
        self.activation_window.geometry("1000x800")
        self.activation_window.transient(self.root)
        
    def show_train_params_dialog(self):
        """显示训练参数设置对话框"""
        dialog = tk.Toplevel(self.root)
        dialog.title("设置训练参数")
        dialog.geometry("600x700")
        dialog.transient(self.root)
        dialog.grab_set()
        
        # 创建滚动条和画布
        canvas = tk.Canvas(dialog)
        scrollbar = ttk.Scrollbar(dialog, orient="vertical", command=canvas.yview)
        scrollable_frame = ttk.Frame(canvas)
        
        scrollable_frame.bind(
            "<Configure>",
            lambda e: canvas.configure(scrollregion=canvas.bbox("all"))
        )
        
        canvas.create_window((0, 0), window=scrollable_frame, anchor="nw")
        canvas.configure(yscrollcommand=scrollbar.set)
        
        canvas.pack(side="left", fill="both", expand=True)
        scrollbar.pack(side="right", fill="y")
        
        # 模型参数组
        model_frame = ttk.LabelFrame(scrollable_frame, text="模型超参数")
        model_frame.pack(fill=tk.X, padx=10, pady=10)
        
        # 卷积层数
        ttk.Label(model_frame, text="卷积层数:").grid(row=0, column=0, padx=10, pady=5, sticky=tk.W)
        num_layers_var = tk.StringVar(value=str(self.train_params['num_conv_layers']))
        ttk.Entry(model_frame, textvariable=num_layers_var, width=10).grid(row=0, column=1, padx=10, pady=5)
        ttk.Label(model_frame, text="(1-5)", font=('Arial', 8, 'italic')).grid(row=0, column=2, padx=5, pady=5, sticky=tk.W)
        
        # 基础通道数
        ttk.Label(model_frame, text="基础通道数:").grid(row=1, column=0, padx=10, pady=5, sticky=tk.W)
        base_channels_var = tk.StringVar(value=str(self.train_params['base_channels']))
        ttk.Entry(model_frame, textvariable=base_channels_var, width=10).grid(row=1, column=1, padx=10, pady=5)
        ttk.Label(model_frame, text="(16, 32, 64等)", font=('Arial', 8, 'italic')).grid(row=1, column=2, padx=5, pady=5, sticky=tk.W)
        
        # LeakyReLU的alpha值
        ttk.Label(model_frame, text="LeakyReLU alpha:").grid(row=2, column=0, padx=10, pady=5, sticky=tk.W)
        alpha_var = tk.StringVar(value=str(self.train_params['leaky_relu_alpha']))
        ttk.Entry(model_frame, textvariable=alpha_var, width=10).grid(row=2, column=1, padx=10, pady=5)
        
        # Dropout率
        ttk.Label(model_frame, text="Dropout率:").grid(row=3, column=0, padx=10, pady=5, sticky=tk.W)
        dropout_var = tk.StringVar(value=str(self.train_params['dropout_rate']))
        ttk.Entry(model_frame, textvariable=dropout_var, width=10).grid(row=3, column=1, padx=10, pady=5)
        ttk.Label(model_frame, text="(0-1)", font=('Arial', 8, 'italic')).grid(row=3, column=2, padx=5, pady=5, sticky=tk.W)
        
        # 批归一化
        use_bn_var = tk.BooleanVar(value=self.train_params['use_batch_norm'])
        ttk.Checkbutton(model_frame, text="使用批归一化", variable=use_bn_var).grid(row=4, column=0, columnspan=3, padx=10, pady=5, sticky=tk.W)
        
        # 全局最大池化
        use_gmp_var = tk.BooleanVar(value=self.train_params['use_gmp'])
        ttk.Checkbutton(model_frame, text="使用全局最大池化", variable=use_gmp_var).grid(row=5, column=0, columnspan=3, padx=10, pady=5, sticky=tk.W)
        
        # MLP路径配置
        ttk.Separator(model_frame, orient='horizontal').grid(row=6, column=0, columnspan=3, sticky=tk.EW, pady=5)
        
        use_mlp_var = tk.BooleanVar(value=self.train_params.get('use_mlp', True))
        ttk.Checkbutton(model_frame, text="使用MLP路径（与CNN融合）", variable=use_mlp_var).grid(row=7, column=0, columnspan=3, padx=10, pady=5, sticky=tk.W)
        
        # MLP隐藏层配置
        ttk.Label(model_frame, text="MLP隐藏层维度（逗号分隔）:").grid(row=8, column=0, padx=10, pady=5, sticky=tk.W)
        mlp_hidden_dims_var = tk.StringVar(value=",".join(map(str, self.train_params.get('mlp_hidden_dims', (256, 128)))))
        ttk.Entry(model_frame, textvariable=mlp_hidden_dims_var, width=20).grid(row=8, column=1, padx=10, pady=5)
        ttk.Label(model_frame, text="例如: 256,128", font=('Arial', 8, 'italic')).grid(row=8, column=2, padx=5, pady=5, sticky=tk.W)
        
        # 训练参数组
        train_frame = ttk.LabelFrame(scrollable_frame, text="训练参数")
        train_frame.pack(fill=tk.X, padx=10, pady=10)
        
        # 批大小
        ttk.Label(train_frame, text="批大小:").grid(row=0, column=0, padx=10, pady=5, sticky=tk.W)
        batch_size_var = tk.StringVar(value=str(self.train_params['batch_size']))
        ttk.Entry(train_frame, textvariable=batch_size_var, width=10).grid(row=0, column=1, padx=10, pady=5)
        
        # 学习率
        ttk.Label(train_frame, text="学习率:").grid(row=1, column=0, padx=10, pady=5, sticky=tk.W)
        lr_var = tk.StringVar(value=str(self.train_params['learning_rate']))
        ttk.Entry(train_frame, textvariable=lr_var, width=10).grid(row=1, column=1, padx=10, pady=5)
        
        # 轮数
        ttk.Label(train_frame, text="轮数:").grid(row=2, column=0, padx=10, pady=5, sticky=tk.W)
        epochs_var = tk.StringVar(value=str(self.train_params['epochs']))
        ttk.Entry(train_frame, textvariable=epochs_var, width=10).grid(row=2, column=1, padx=10, pady=5)
        
        # 权重衰减
        ttk.Label(train_frame, text="权重衰减:").grid(row=3, column=0, padx=10, pady=5, sticky=tk.W)
        weight_decay_var = tk.StringVar(value=str(self.train_config['weight_decay']))
        ttk.Entry(train_frame, textvariable=weight_decay_var, width=10).grid(row=3, column=1, padx=10, pady=5)
        
        # 梯度裁剪
        ttk.Label(train_frame, text="梯度裁剪:").grid(row=4, column=0, padx=10, pady=5, sticky=tk.W)
        gradient_clip_var = tk.StringVar(value=str(self.train_config['gradient_clip']))
        ttk.Entry(train_frame, textvariable=gradient_clip_var, width=10).grid(row=4, column=1, padx=10, pady=5)
        
        # 早停耐心值
        ttk.Label(train_frame, text="早停耐心值:").grid(row=5, column=0, padx=10, pady=5, sticky=tk.W)
        patience_var = tk.StringVar(value=str(self.train_config['early_stopping_patience']))
        ttk.Entry(train_frame, textvariable=patience_var, width=10).grid(row=5, column=1, padx=10, pady=5)
        
        # 早停启用
        early_stopping_var = tk.BooleanVar(value=self.train_config['early_stopping'])
        ttk.Checkbutton(train_frame, text="启用早停", variable=early_stopping_var).grid(row=6, column=0, columnspan=2, padx=10, pady=5, sticky=tk.W)
        
        # 结果显示参数
        ttk.Separator(train_frame, orient='horizontal').grid(row=7, column=0, columnspan=3, sticky=tk.EW, pady=5)
        
        ttk.Label(train_frame, text="显示预测结果数量:").grid(row=8, column=0, padx=10, pady=5, sticky=tk.W)
        top_n_var = tk.StringVar(value=str(self.train_config.get('top_n_predictions', 3)))
        ttk.Entry(train_frame, textvariable=top_n_var, width=10).grid(row=8, column=1, padx=10, pady=5)
        ttk.Label(train_frame, text="(1-10)", font=('Arial', 8, 'italic')).grid(row=8, column=2, padx=5, pady=5, sticky=tk.W)
        
        # 学习率调度器
        ttk.Label(train_frame, text="学习率调度器:").grid(row=9, column=0, padx=10, pady=5, sticky=tk.W)
        lr_scheduler_var = tk.StringVar(value=self.train_config['lr_scheduler'])
        
        # 数据增强参数组
        aug_frame = ttk.LabelFrame(scrollable_frame, text="数据增强设置")
        aug_frame.pack(fill=tk.X, padx=10, pady=10)

        # 数据增强级别
        ttk.Label(aug_frame, text="数据增强级别:").grid(row=0, column=0, padx=10, pady=5, sticky=tk.W)
        augmentation_level_var = tk.StringVar(value=self.train_config.get('augmentation_level', 'medium'))
        augmentation_combobox = ttk.Combobox(aug_frame, textvariable=augmentation_level_var, 
                                        values=['medium', 'heavy'], width=10)
        augmentation_combobox.grid(row=0, column=1, padx=10, pady=5)

        # 使用数据增强
        use_augmentation_var = tk.BooleanVar(value=self.train_config.get('use_augmentation', True))
        ttk.Checkbutton(aug_frame, text="启用数据增强", variable=use_augmentation_var).grid(row=1, column=0, columnspan=2, padx=10, pady=5, sticky=tk.W)

        # 确保scheduler_params存在
        if 'scheduler_params' not in self.train_config:
            self.train_config['scheduler_params'] = {
                'onecycle': {'max_lr_multiplier': 25.0, 'pct_start': 0.3},
                'cosine': {'T_max': 10, 'min_lr_multiplier': 0.01},
                'step': {'step_size': 5, 'gamma': 0.5}
            }
        
        # 创建调度器下拉框
        scheduler_combobox = ttk.Combobox(train_frame, textvariable=lr_scheduler_var, values=['onecycle', 'cosine', 'step'], width=8)
        scheduler_combobox.grid(row=9, column=1, padx=10, pady=5)
        
        # 创建调度器参数框架
        scheduler_params_frame = ttk.LabelFrame(train_frame, text="调度器参数")
        scheduler_params_frame.grid(row=10, column=0, columnspan=3, padx=10, pady=5, sticky=tk.W+tk.E)
        
        # OneCycle参数
        onecycle_frame = ttk.Frame(scheduler_params_frame)
        onecycle_frame.pack(fill=tk.X, padx=10, pady=5)
        
        ttk.Label(onecycle_frame, text="Max LR倍率:").grid(row=0, column=0, padx=10, pady=2, sticky=tk.W)
        onecycle_max_lr_var = tk.StringVar(value=str(self.train_config['scheduler_params']['onecycle']['max_lr_multiplier']))
        ttk.Entry(onecycle_frame, textvariable=onecycle_max_lr_var, width=10).grid(row=0, column=1, padx=10, pady=2)
        
        ttk.Label(onecycle_frame, text="Warmup比例:").grid(row=1, column=0, padx=10, pady=2, sticky=tk.W)
        onecycle_pct_start_var = tk.StringVar(value=str(self.train_config['scheduler_params']['onecycle']['pct_start']))
        ttk.Entry(onecycle_frame, textvariable=onecycle_pct_start_var, width=10).grid(row=1, column=1, padx=10, pady=2)
        
        # Cosine参数
        cosine_frame = ttk.Frame(scheduler_params_frame)
        cosine_frame.pack(fill=tk.X, padx=10, pady=5)
        
        ttk.Label(cosine_frame, text="退火周期:").grid(row=0, column=0, padx=10, pady=2, sticky=tk.W)
        cosine_tmax_var = tk.StringVar(value=str(self.train_config['scheduler_params']['cosine']['T_max']))
        ttk.Entry(cosine_frame, textvariable=cosine_tmax_var, width=10).grid(row=0, column=1, padx=10, pady=2)
        
        ttk.Label(cosine_frame, text="最小LR倍率:").grid(row=1, column=0, padx=10, pady=2, sticky=tk.W)
        cosine_min_lr_var = tk.StringVar(value=str(self.train_config['scheduler_params']['cosine']['min_lr_multiplier']))
        ttk.Entry(cosine_frame, textvariable=cosine_min_lr_var, width=10).grid(row=1, column=1, padx=10, pady=2)
        
        # Step参数
        step_frame = ttk.Frame(scheduler_params_frame)
        step_frame.pack(fill=tk.X, padx=10, pady=5)
        
        ttk.Label(step_frame, text="步长:").grid(row=0, column=0, padx=10, pady=2, sticky=tk.W)
        step_size_var = tk.StringVar(value=str(self.train_config['scheduler_params']['step']['step_size']))
        ttk.Entry(step_frame, textvariable=step_size_var, width=10).grid(row=0, column=1, padx=10, pady=2)
        
        ttk.Label(step_frame, text="衰减因子:").grid(row=1, column=0, padx=10, pady=2, sticky=tk.W)
        step_gamma_var = tk.StringVar(value=str(self.train_config['scheduler_params']['step']['gamma']))
        ttk.Entry(step_frame, textvariable=step_gamma_var, width=10).grid(row=1, column=1, padx=10, pady=2)
        
        # 控制参数框架显示的函数
        def update_scheduler_params_visibility(event=None):
            scheduler_type = lr_scheduler_var.get()
            onecycle_frame.pack_forget()
            cosine_frame.pack_forget()
            step_frame.pack_forget()
            
            if scheduler_type == 'onecycle':
                onecycle_frame.pack(fill=tk.X, padx=10, pady=5)
            elif scheduler_type == 'cosine':
                cosine_frame.pack(fill=tk.X, padx=10, pady=5)
            elif scheduler_type == 'step':
                step_frame.pack(fill=tk.X, padx=10, pady=5)
        
        # 绑定事件和初始化显示
        scheduler_combobox.bind("<<ComboboxSelected>>", update_scheduler_params_visibility)
        update_scheduler_params_visibility()
        
        def save_params():
            try:
                # 保存模型参数
                self.train_params['num_conv_layers'] = int(num_layers_var.get())
                self.train_params['base_channels'] = int(base_channels_var.get())
                self.train_params['leaky_relu_alpha'] = float(alpha_var.get())
                self.train_params['dropout_rate'] = float(dropout_var.get())
                self.train_params['use_batch_norm'] = use_bn_var.get()
                self.train_params['use_gmp'] = use_gmp_var.get()
                
                # 保存MLP相关参数
                self.train_params['use_mlp'] = use_mlp_var.get()
                # 解析MLP隐藏层维度
                try:
                    mlp_dims_str = mlp_hidden_dims_var.get().strip()
                    if mlp_dims_str:
                        mlp_dims = tuple(map(int, mlp_dims_str.split(',')))
                        for dim in mlp_dims:
                            if dim <= 0 or dim > 1024:
                                raise ValueError("MLP隐藏层维度必须在1-1024之间")
                        self.train_params['mlp_hidden_dims'] = mlp_dims
                    else:
                        raise ValueError("MLP隐藏层维度不能为空")
                except Exception as e:
                    raise ValueError(f"MLP隐藏层维度格式错误: {str(e)}")
                
                # 保存训练参数
                self.train_params['batch_size'] = int(batch_size_var.get())
                self.train_params['learning_rate'] = float(lr_var.get())
                self.train_params['epochs'] = int(epochs_var.get())
                
                # 保存配置参数
                self.train_config['batch_size'] = self.train_params['batch_size']
                self.train_config['learning_rate'] = self.train_params['learning_rate']
                self.train_config['epochs'] = self.train_params['epochs']
                self.train_config['weight_decay'] = float(weight_decay_var.get())
                self.train_config['gradient_clip'] = float(gradient_clip_var.get())
                self.train_config['early_stopping_patience'] = int(patience_var.get())
                self.train_config['early_stopping'] = early_stopping_var.get()
                self.train_config['lr_scheduler'] = lr_scheduler_var.get()
                self.train_config['top_n_predictions'] = int(top_n_var.get())
                
                # 保存调度器参数
                self.train_config['scheduler_params']['onecycle']['max_lr_multiplier'] = float(onecycle_max_lr_var.get())
                self.train_config['scheduler_params']['onecycle']['pct_start'] = float(onecycle_pct_start_var.get())
                self.train_config['scheduler_params']['cosine']['T_max'] = int(cosine_tmax_var.get())
                self.train_config['scheduler_params']['cosine']['min_lr_multiplier'] = float(cosine_min_lr_var.get())
                self.train_config['scheduler_params']['step']['step_size'] = int(step_size_var.get())
                self.train_config['scheduler_params']['step']['gamma'] = float(step_gamma_var.get())
                
                # 保存数据增强参数
                self.train_config['augmentation_level'] = augmentation_level_var.get()
                self.train_config['use_augmentation'] = use_augmentation_var.get()

                # 验证调度器参数
                if self.train_config['scheduler_params']['onecycle']['max_lr_multiplier'] <= 0:
                    raise ValueError("OneCycle最大学习率倍率必须为正数")
                if not 0 < self.train_config['scheduler_params']['onecycle']['pct_start'] < 1:
                    raise ValueError("OneCycle pct_start必须在0到1之间")
                if self.train_config['scheduler_params']['cosine']['T_max'] <= 0:
                    raise ValueError("Cosine退火周期必须为正数")
                if self.train_config['scheduler_params']['cosine']['min_lr_multiplier'] <= 0:
                    raise ValueError("Cosine最小学习率倍率必须为正数")
                if self.train_config['scheduler_params']['step']['step_size'] <= 0:
                    raise ValueError("Step步长必须为正数")
                if not 0 < self.train_config['scheduler_params']['step']['gamma'] < 1:
                    raise ValueError("Step衰减因子必须在0到1之间")
                
                # 验证参数
                if self.train_params['batch_size'] <= 0 or self.train_params['epochs'] <= 0:
                    raise ValueError("批大小和轮数必须为正数")
                if self.train_params['learning_rate'] <= 0:
                    raise ValueError("学习率必须为正数")
                if self.train_params['base_channels'] <= 0 or self.train_params['base_channels'] > 128:
                    raise ValueError("基础通道数必须在1-128之间")
                if self.train_params['leaky_relu_alpha'] < 0:
                    raise ValueError("LeakyReLU alpha不能为负数")
                if self.train_params['dropout_rate'] < 0 or self.train_params['dropout_rate'] >= 1:
                    raise ValueError("Dropout率必须在0到1之间")
                if self.train_params['num_conv_layers'] < 1 or self.train_params['num_conv_layers'] > 5:
                    raise ValueError("卷积层数必须在1-5之间")
                if self.train_config['top_n_predictions'] < 1 or self.train_config['top_n_predictions'] > 10:
                    raise ValueError("预测结果显示数量必须在1-10之间")
                
                # 更新top_n_predictions
                self.top_n_predictions = self.train_config['top_n_predictions']
                
                dialog.destroy()
                messagebox.showinfo("成功", "参数保存成功!")
            except ValueError as e:
                messagebox.showerror("错误", f"参数无效: {str(e)}")
        
        # 保存按钮
        ttk.Button(scrollable_frame, text="保存", command=save_params).pack(pady=20)
        # 取消按钮
        ttk.Button(scrollable_frame, text="取消", command=dialog.destroy).pack(pady=5)
    
    def start_training(self):
        """开始训练模型"""
        # 重置训练数据和图表
        def reset_training_data():
            # 清空训练历史数据
            self.epochs = []
            self.train_losses = []
            self.test_losses = []
            self.train_accs = []
            self.test_accs = []
            
            # 重置图表数据
            if hasattr(self, 'loss_visualizer') and self.loss_visualizer:
                self.loss_visualizer.clear()
                self.acc_visualizer.clear()
        
        # 在主线程中重置训练数据
        self.root.after(0, reset_training_data)
        
        # 使用自定义参数重新加载数据
        try:
            # 根据augmentation_level确定是否使用高级数据增强
            advanced_augmentation = self.train_config['augmentation_level'] in ['medium', 'heavy']
            
            # 加载数据集
            self.data_loader.load_datasets(
                validation_split=0.1,
                advanced_augmentation=advanced_augmentation
            )
            
            # 创建数据加载器
            self.train_loader, self.val_loader, self.test_loader = self.data_loader.create_dataloaders(
                batch_size=self.train_params['batch_size'],
                num_workers=4,
                shuffle_train=True,
                pin_memory=True,
                persistent_workers=True
            )
        except Exception as e:
            messagebox.showerror("错误", f"数据加载失败: {str(e)}")
            return
        
        # 创建模型 - 使用自定义参数
        try:
            # MNIST是0-9数字识别，类别数固定为10
            print(f"创建模型，参数: {self.train_params}")
            # 先在CPU上创建模型
            self.model = CNN(
                num_classes=self.train_params['num_classes'],
                num_conv_layers=self.train_params['num_conv_layers'],
                base_channels=self.train_params['base_channels'],
                alpha=self.train_params['leaky_relu_alpha'],
                dropout_rate=self.train_params['dropout_rate'],
                use_batch_norm=self.train_params['use_batch_norm'],
                use_gmp=self.train_params['use_gmp'],
                use_mlp=self.train_params.get('use_mlp', True),
                mlp_hidden_dims=self.train_params.get('mlp_hidden_dims', (256, 128))
            )
            
            # 强制重新初始化权重
            self.model._initialize_weights()
            
            # 将模型移至正确设备
            self.model = self.model.to(self.device)
            
            # 确保scheduler_params存在
            if 'scheduler_params' not in self.train_config:
                self.train_config['scheduler_params'] = {
                    'onecycle': {'max_lr_multiplier': 25.0, 'pct_start': 0.3},
                    'cosine': {'T_max': 10, 'min_lr_multiplier': 0.01},
                    'step': {'step_size': 5, 'gamma': 0.5}
                }
            
            # 创建训练器
            self.trainer = Trainer(
                self.model, 
                self.train_loader, 
                self.test_loader, 
                device=self.device,
                lr=self.train_params['learning_rate'],
                weight_decay=self.train_config['weight_decay'],
                betas=self.train_config['betas'],
                early_stopping_patience=self.train_config['early_stopping_patience'],
                gradient_clip=self.train_config['gradient_clip'],
                lr_scheduler_type=self.train_config['lr_scheduler']
            )
            
        except Exception as e:
            messagebox.showerror("错误", f"模型创建失败: {str(e)}")
            return
        
        # 设置回调函数更新可视化
        def on_epoch_end(trainer, epoch):
            # 确保应用程序仍在运行且没有被清理
            if not self.running or self._cleaning_up:
                return
                
            # 确保GUI更新在主线程中进行
            def update_gui():
                if not self.running or self._cleaning_up:
                    return
                       
                try:
                    self.epochs.append(epoch + 1)
                    self.train_losses.append(trainer.train_loss_history[-1])
                    self.test_losses.append(trainer.test_loss_history[-1])
                    self.train_accs.append(trainer.train_acc_history[-1])
                    self.test_accs.append(trainer.test_acc_history[-1])
                    
                    # 更新图表
                    history_data = {
                        'epoch': self.epochs[-1],
                        'train_loss': self.train_losses[-1],
                        'test_loss': self.test_losses[-1],
                        'train_acc': self.train_accs[-1],
                        'test_acc': self.test_accs[-1]
                    }
                    if hasattr(self, 'loss_visualizer') and self.loss_visualizer:
                        self.loss_visualizer.update(history_data)
                        self.acc_visualizer.update(history_data)
                    
                    # 更新进度
                    self.progress['value'] = ((epoch + 1) / self.train_params['epochs']) * 100
                except Exception as e:
                    # 忽略清理过程中可能发生的错误
                    if self.running and not self._cleaning_up:
                        print(f"更新GUI时出错: {str(e)}")
            
            # 使用after方法确保GUI更新在主线程中进行
            self.root.after(0, update_gui)
        
        self.trainer.on_epoch_end(on_epoch_end)
        
        # 在单独线程中训练
        def train_thread():
            # 使用after方法更新GUI状态
            self.root.after(0, lambda: self.train_btn.config(state=tk.DISABLED))
            self.root.after(0, lambda: self.cancel_train_btn.config(state=tk.NORMAL))
            self.root.after(0, lambda: self.progress.start(10))
            
            try:
                # 进行训练
                self.trainer.fit(
                    epochs=self.train_params['epochs'],
                    early_stopping=self.train_config['early_stopping'],
                    patience=self.train_config['early_stopping_patience'],
                    scheduler_params=self.train_config['scheduler_params']
                )
                
                # 仅在正常完成训练时显示消息
                if self.running and not self.training_cancelled:
                    best_acc = self.trainer.get_best_accuracy() if hasattr(self.trainer, 'get_best_accuracy') else 0
                    self.root.after(0, lambda: messagebox.showinfo("成功", f"模型训练完成!\n最佳测试准确率: {best_acc:.2f}%"))
                else:
                    self.root.after(0, lambda: messagebox.showinfo("信息", "训练已中断"))
            except Exception as e:
                # 错误消息也在主线程中显示
                if self.running and not self._cleaning_up:
                    self.root.after(0, lambda err_msg=str(e): messagebox.showerror("错误", f"训练失败: {err_msg}"))
            finally:
                # 恢复GUI状态
                self.root.after(0, lambda: self.train_btn.config(state=tk.NORMAL))
                self.root.after(0, lambda: self.cancel_train_btn.config(state=tk.DISABLED))
                self.root.after(0, lambda: self.progress.stop())
        
        # 使用非守护线程，确保训练过程不会被意外终止
        self.train_thread = threading.Thread(target=train_thread, daemon=False)
        self.train_thread.start()
        
    def cancel_training(self):
        """取消训练过程"""
        if messagebox.askyesno("确认", "确定要取消训练吗？"):
            self.training_cancelled = True
            
            # 通知训练器停止训练
            if hasattr(self, 'trainer') and self.trainer:
                self.trainer.stop_training()
                print("[应用程序] 已发送停止训练信号")
            
            # 立即更新UI状态
            self.train_btn.config(state=tk.NORMAL)
            self.cancel_train_btn.config(state=tk.DISABLED)
            self.progress.stop()
            
            # 等待训练线程结束
            def wait_for_thread():
                if hasattr(self, 'train_thread') and self.train_thread.is_alive():
                    # 线程仍在运行，继续等待
                    self.root.after(100, wait_for_thread)
                else:
                    # 线程已停止，重置标志
                    self.training_cancelled = False
                    print("[应用程序] 训练线程已完全停止")
            
            # 开始等待线程结束
            wait_for_thread()
    
    def save_model(self):
        """保存模型"""
        if self.model is None:
            messagebox.showwarning("警告", "没有模型可保存!")
            return
        
        file_path = filedialog.asksaveasfilename(
            defaultextension=".pth",
            filetypes=[("PyTorch模型", "*.pth"), ("所有文件", "*.*")]
        )
        
        if file_path:
            try:
                # 保存模型权重和配置参数
                checkpoint = {
                    'model_state_dict': self.model.state_dict(),
                    'train_params': self.train_params.copy(),  # 保存创建模型时的参数配置
                    'model_info': f"CNN_{self.train_params['num_conv_layers']}layers_{self.train_params['base_channels']}channels"
                }
                torch.save(checkpoint, file_path)
                messagebox.showinfo("成功", "模型保存成功!")
            except Exception as e:
                messagebox.showerror("错误", f"保存失败: {str(e)}")
    
    def load_model(self):
        """加载模型"""
        file_path = filedialog.askopenfilename(
            filetypes=[("PyTorch模型", "*.pth"), ("所有文件", "*.*")]
        )
        
        if file_path:
            try:
                # 先加载检查点以获取模型配置
                checkpoint = torch.load(file_path, map_location=self.device)
                
                # 检查是否是包含模型配置的完整检查点
                if isinstance(checkpoint, dict):
                    # 如果检查点中包含模型配置参数，则使用这些参数创建模型
                    if 'train_params' in checkpoint:
                        model_params = checkpoint['train_params']
                        print(f"从检查点加载模型配置: {model_params}")
                        # 保存这些参数到当前配置
                        self.train_params.update(model_params)
                        
                        # 使用保存的参数创建模型
                        self.model = CNN(
                            num_classes=model_params['num_classes'],
                            num_conv_layers=model_params['num_conv_layers'],
                            base_channels=model_params['base_channels'],
                            alpha=model_params['leaky_relu_alpha'],
                            dropout_rate=model_params['dropout_rate'],
                            use_batch_norm=model_params['use_batch_norm'],
                            use_gmp=model_params['use_gmp'],
                            use_mlp=model_params.get('use_mlp', True),
                            mlp_hidden_dims=model_params.get('mlp_hidden_dims', (256, 128))
                        ).to(self.device)
                        
                        # 加载模型权重
                        if 'model_state_dict' in checkpoint:
                            self.model.load_state_dict(checkpoint['model_state_dict'])
                            print("从完整检查点加载模型权重")
                        else:
                            self.model.load_state_dict(checkpoint)
                            print("加载模型权重成功")
                    # 如果是简单的检查点格式
                    elif 'model_state_dict' in checkpoint:
                        # 使用当前配置创建模型
                        print(f"使用当前配置创建模型: {self.train_params}")
                        self.model = CNN(
                            num_classes=self.train_params['num_classes'],
                            num_conv_layers=self.train_params['num_conv_layers'],
                            base_channels=self.train_params['base_channels'],
                            alpha=self.train_params['leaky_relu_alpha'],
                            dropout_rate=self.train_params['dropout_rate'],
                            use_batch_norm=self.train_params['use_batch_norm'],
                            use_gmp=self.train_params['use_gmp'],
                            use_mlp=self.train_params.get('use_mlp', True),
                            mlp_hidden_dims=self.train_params.get('mlp_hidden_dims', (256, 128))
                        ).to(self.device)
                        
                        # 加载模型权重
                        self.model.load_state_dict(checkpoint['model_state_dict'])
                        print("从包含state_dict的检查点加载模型权重")
                    else:
                        # 尝试作为简单的state_dict加载
                        print(f"使用当前配置创建模型: {self.train_params}")
                        self.model = CNN(
                            num_classes=self.train_params['num_classes'],
                            num_conv_layers=self.train_params['num_conv_layers'],
                            base_channels=self.train_params['base_channels'],
                            alpha=self.train_params['leaky_relu_alpha'],
                            dropout_rate=self.train_params['dropout_rate'],
                            use_batch_norm=self.train_params['use_batch_norm'],
                            use_gmp=self.train_params['use_gmp'],
                            use_mlp=self.train_params.get('use_mlp', True),
                            mlp_hidden_dims=self.train_params.get('mlp_hidden_dims', (256, 128))
                        ).to(self.device)
                        
                        # 加载模型权重
                        self.model.load_state_dict(checkpoint)
                        print("作为简单的state_dict加载成功")
                else:
                    # 对于旧格式的模型（只有state_dict）
                    print(f"使用当前配置创建模型: {self.train_params}")
                    self.model = CNN(
                        num_classes=self.train_params['num_classes'],
                        num_conv_layers=self.train_params['num_conv_layers'],
                        base_channels=self.train_params['base_channels'],
                        alpha=self.train_params['leaky_relu_alpha'],
                        dropout_rate=self.train_params['dropout_rate'],
                        use_batch_norm=self.train_params['use_batch_norm'],
                        use_gmp=self.train_params['use_gmp'],
                        use_mlp=self.train_params.get('use_mlp', True),
                        mlp_hidden_dims=self.train_params.get('mlp_hidden_dims', (256, 128))
                    ).to(self.device)
                    
                    # 加载模型权重
                    self.model.load_state_dict(checkpoint)
                    print("加载简单模型权重成功")
                
                # 创建训练器
                self.trainer = Trainer(self.model, self.train_loader, self.test_loader, 
                                             device=self.device, lr=self.train_params['learning_rate'])
                messagebox.showinfo("成功", "模型加载成功!")
            except Exception as e:
                messagebox.showerror("错误", f"加载失败: {str(e)}")
                print(f"模型加载错误: {str(e)}")

def main():
    import threading
    import sys
    import matplotlib.pyplot as plt
    
    root = tk.Tk()
    app = MNISTApp(root)
    
    # 确保窗口关闭时正确终止线程
    def on_closing():
        print("正在关闭应用程序，清理线程...")
        
        app.running = False
        app._cleaning_up = True
        
        # 如果有训练器，发送停止信号
        if hasattr(app, 'trainer') and app.trainer:
            app.trainer.stop_training()
        
        all_threads = threading.enumerate()
        print(f"当前活跃线程数: {len(all_threads)}")
        
        try:
            plt.close('all')
            # 断开matplotlib的事件连接
            if hasattr(app, 'fig') and app.fig:
                app.fig.clf()
            if hasattr(app, 'canvas_viz') and app.canvas_viz:
                app.canvas_viz.stop_event_loop()
                app.canvas_viz = None
        except Exception as e:
            print(f"清理matplotlib资源时出错: {str(e)}")
        
        # 等待训练线程结束，但不阻塞主线程过长时间
        def wait_for_threads():
            # 等待训练线程结束
            if app.train_thread and app.train_thread.is_alive():
                print("等待训练线程结束...")
                app.train_thread.join(timeout=1.0)
                if app.train_thread.is_alive():
                    print("警告: 训练线程未能在指定时间内结束")
                
            # 清理识别线程
            for thread in app.recognize_threads[:]:
                if thread.is_alive():
                    print(f"等待识别线程 {thread.name} 结束...")
                    thread.join(timeout=0.3)
                    if thread.is_alive():
                        print(f"警告: 识别线程 {thread.name} 未能在指定时间内结束")
                
            # 强制关闭窗口并退出
            root.after(0, force_quit)
        
        def force_quit():
            try:
                # 先尝试正常销毁窗口
                root.destroy()
                # 强制退出Python进程，确保所有线程终止
                sys.exit(0)
            except Exception as e:
                print(f"销毁窗口时出错: {str(e)}")
                # 最极端情况，强制退出
                os._exit(0)
        
        # 先隐藏窗口，避免进一步的用户交互
        root.withdraw()
        
        # 然后等待线程结束并最终退出
        root.after(100, wait_for_threads)
    
    root.protocol("WM_DELETE_WINDOW", on_closing)
    root.mainloop()

if __name__ == "__main__":
    main()